import logging
import time

from PIL import ImageDraw, ImageFont
import sqlite3

import Config

# 定义一个颜色列表，每个颜色对应一个类别
class_colors = [
    (255, 0, 0),  # 红色
    (255, 255, 0),  # 黄色
    # (255, 255, 255),  # 白色
    (0, 0, 255),  # 蓝色
    (0, 255, 0),  # 绿色
    (255, 0, 255),  # 品红
    (0, 255, 255),  # 青色
]


def detect(model, img):
    results = model(img)
    return results


def draw_img(model, results, img):
    draw = ImageDraw.Draw(img)
    # 指定字体和字体大小
    font_path = Config.ttc  # 确保这是正确的字体路径
    font_size = 50
    font = ImageFont.truetype(font_path, font_size)

    for result in results:
        for det in result.boxes:  # 遍历每个检测到的对象
            x1, y1, x2, y2, conf, cls = int(det.xyxy[0][0]), int(det.xyxy[0][1]), int(det.xyxy[0][2]), int(
                det.xyxy[0][3]), \
                det.conf[0], int(det.cls[0])

            # 根据类别索引选择颜色
            color = class_colors[cls % len(class_colors)]

            # 绘制边界框
            draw.rectangle((x1, y1, x2, y2), outline=color, width=2)

            # 获取类别名称
            class_name = model.names[cls]

            # 绘制类别名称
            draw.text((x1, y1 - 20), class_name, font=font, fill=color)
    return img


def save_data(mission, path):
    conn = sqlite3.connect(Config.databases)
    cursor = conn.cursor()
    try:
        time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        data = [mission, path, time_str]
        cursor.execute("INSERT INTO history (type, path, time) VALUES (?, ?, ?)", data)
        conn.commit()
        logging.info('保存成功')
    except sqlite3.OperationalError and sqlite3.Error as error:
        logging.error('保存失败', exc_info=True)